# iris数据集 KNN算法
import matplotlib.pyplot as plt
from sklearn.neighbors import KNeighborsClassifier
from sklearn import datasets
import sklearn
import numpy as np
from sklearn.model_selection import train_test_split  # 更新

np.random.seed(0)
iris = datasets.load_iris()
x_data = iris.data
y_data = iris.target

# 随机抽取30%的测试集
x_train, x_test, y_train, y_test = sklearn.model_selection.train_test_split(
    x_data, y_data, test_size=0.3)

print(x_train.shape,y_train.shape,x_test.shape,y_test.shape)

knn = KNeighborsClassifier(n_neighbors=3) # KNN
knn.fit(x_train, y_train)

score=knn.score(x_train,y_train)
print("训练集的准确率：", score)
score=knn.score(x_test,y_test)
print("测试集的准确率：", score)

y_pred = knn.predict(x_test)
print(y_pred)